/*-------------------------------------------------------------------------
 *
 * partitionfuncs.c
 *	  Functions for accessing partition-related metadata
 *
 * Portions Copyright (c) 1996-2018, PostgreSQL Global Development Group
 * Portions Copyright (c) 1994, Regents of the University of California
 *
 * This source code file contains modifications made by THL A29 Limited ("Tencent Modifications").
 * All Tencent Modifications are Copyright (C) 2023 THL A29 Limited.
 *
 * IDENTIFICATION
 *	  src/backend/utils/adt/partitionfuncs.c
 *
 *-------------------------------------------------------------------------
 */

#include "postgres.h"

#include "access/htup_details.h"
#include "catalog/partition.h"
#include "catalog/pg_class.h"
#include "catalog/pg_inherits.h"
#include "catalog/pg_inherits_fn.h"
#include "catalog/pg_type.h"
#include "funcapi.h"
#include "utils/fmgrprotos.h"
#include "utils/lsyscache.h"
#include "utils/syscache.h"

/*
* Checks if a given relation can be part of a partition tree.  Returns
* false if the relation cannot be processed, in which case it is up to
* the caller to decide what to do, by either raising an error or doing
* something else.
*/
static bool
check_rel_can_be_partition(Oid relid)
{
    char        relkind;
    bool        relispartition;

    /* Check if relation exists */
    if (!SearchSysCacheExists1(RELOID, ObjectIdGetDatum(relid)))
        return false;

    relkind = get_rel_relkind(relid);
    relispartition = get_rel_relispartition(relid);

    /* Only allow relation types that can appear in partition trees. */
    if (!relispartition &&
        relkind != RELKIND_PARTITIONED_TABLE &&
        relkind != RELKIND_PARTITIONED_INDEX)
        return false;

    return true;
}

/*
 * pg_partition_tree
 *
 * Produce a view with one row per member of a partition tree, beginning
 * from the top-most parent given by the caller.  This gives information
 * about each partition, its immediate partitioned parent, if it is
 * a leaf partition and its level in the hierarchy.
 */
Datum
pg_partition_tree(PG_FUNCTION_ARGS)
{
#define PG_PARTITION_TREE_COLS	4
	Oid			rootrelid = PG_GETARG_OID(0);
	FuncCallContext *funcctx;
	ListCell  **next;

	/* stuff done only on the first call of the function */
	if (SRF_IS_FIRSTCALL())
	{
		MemoryContext oldcxt;
		TupleDesc	tupdesc;
		List	   *partitions;

		/* create a function context for cross-call persistence */
		funcctx = SRF_FIRSTCALL_INIT();

		if (!check_rel_can_be_partition(rootrelid))
           SRF_RETURN_DONE(funcctx);

		/* switch to memory context appropriate for multiple function calls */
		oldcxt = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);

		/*
		 * Find all members of inheritance set.  We only need AccessShareLock
		 * on the children for the partition information lookup.
		 */
		partitions = find_all_inheritors(rootrelid, AccessShareLock, NULL);

		tupdesc = CreateTemplateTupleDesc(PG_PARTITION_TREE_COLS, false);
		TupleDescInitEntry(tupdesc, (AttrNumber) 1, "relid",
						   REGCLASSOID, -1, 0);
		TupleDescInitEntry(tupdesc, (AttrNumber) 2, "parentid",
						   REGCLASSOID, -1, 0);
		TupleDescInitEntry(tupdesc, (AttrNumber) 3, "isleaf",
						   BOOLOID, -1, 0);
		TupleDescInitEntry(tupdesc, (AttrNumber) 4, "level",
						   INT4OID, -1, 0);

		funcctx->tuple_desc = BlessTupleDesc(tupdesc);

		/* allocate memory for user context */
		next = (ListCell **) palloc(sizeof(ListCell *));
		*next = list_head(partitions);
		funcctx->user_fctx = (void *) next;

		MemoryContextSwitchTo(oldcxt);
	}

	/* stuff done on every call of the function */
	funcctx = SRF_PERCALL_SETUP();
	next = (ListCell **) funcctx->user_fctx;

	if (*next != NULL)
	{
		Datum		result;
		Datum		values[PG_PARTITION_TREE_COLS];
		bool		nulls[PG_PARTITION_TREE_COLS];
		HeapTuple	tuple;
		Oid			parentid = InvalidOid;
		Oid			relid = lfirst_oid(*next);
		char		relkind = get_rel_relkind(relid);
		int			level = 0;
		List	   *ancestors = get_partition_ancestors(lfirst_oid(*next));
		ListCell   *lc;

		/*
		 * Form tuple with appropriate data.
		 */
		MemSet(nulls, 0, sizeof(nulls));
		MemSet(values, 0, sizeof(values));

		/* relid */
		values[0] = ObjectIdGetDatum(relid);

		/* parentid */
		if (ancestors != NIL)
			parentid = linitial_oid(ancestors);
		if (OidIsValid(parentid))
			values[1] = ObjectIdGetDatum(parentid);
		else
			nulls[1] = true;

		/* isleaf */
		values[2] = BoolGetDatum(relkind != RELKIND_PARTITIONED_TABLE &&
								 relkind != RELKIND_PARTITIONED_INDEX);

		/* level */
		if (relid != rootrelid)
		{
			foreach(lc, ancestors)
			{
				level++;
				if (lfirst_oid(lc) == rootrelid)
					break;
			}
		}
		values[3] = Int32GetDatum(level);

		*next = lnext(*next);

		tuple = heap_form_tuple(funcctx->tuple_desc, values, nulls);
		result = HeapTupleGetDatum(tuple);
		SRF_RETURN_NEXT(funcctx, result);
	}

	/* done when there are no more elements left */
	SRF_RETURN_DONE(funcctx);
}

/*
 * pg_partition_root
 *
 * Returns the top-most parent of the partition tree to which a given
 * relation belongs, or NULL if it's not (or cannot be) part of any
 * partition tree.
 */
Datum
pg_partition_root(PG_FUNCTION_ARGS)
{
   Oid         relid = PG_GETARG_OID(0);
   Oid         rootrelid;
   List       *ancestors;

   if (!check_rel_can_be_partition(relid))
       PG_RETURN_NULL();

   /* fetch the list of ancestors */
   ancestors = get_partition_ancestors(relid);

   /*
    * If the input relation is already the top-most parent, just return
    * itself.
    */
   if (ancestors == NIL)
       PG_RETURN_OID(relid);

   rootrelid = llast_oid(ancestors);
   list_free(ancestors);

   /*
    * "rootrelid" must contain a valid OID, given that the input relation is
    * a valid partition tree member as checked above.
    */
   Assert(OidIsValid(rootrelid));
   PG_RETURN_OID(rootrelid);
}

/*
 * pg_partition_ancestors
 *
 * Produces a view with one row per ancestor of the given partition,
 * including the input relation itself.
 */
Datum
pg_partition_ancestors(PG_FUNCTION_ARGS)
{
   Oid         relid = PG_GETARG_OID(0);
   FuncCallContext *funcctx;
   ListCell  **next;

   if (SRF_IS_FIRSTCALL())
   {
       MemoryContext oldcxt;
       List       *ancestors;

       funcctx = SRF_FIRSTCALL_INIT();

       if (!check_rel_can_be_partition(relid))
           SRF_RETURN_DONE(funcctx);

       oldcxt = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);

       ancestors = get_partition_ancestors(relid);
       ancestors = lcons_oid(relid, ancestors);

       next = (ListCell **) palloc(sizeof(ListCell *));
       *next = list_head(ancestors);
       funcctx->user_fctx = (void *) next;

       MemoryContextSwitchTo(oldcxt);
   }

   funcctx = SRF_PERCALL_SETUP();
   next = (ListCell **) funcctx->user_fctx;

   if (*next != NULL)
   {
       Oid         relid = lfirst_oid(*next);

       *next = lnext(*next);
       SRF_RETURN_NEXT(funcctx, ObjectIdGetDatum(relid));
   }

   SRF_RETURN_DONE(funcctx);
}
